import asyncio
import json
import random
import string
import time
from fastapi import FastAPI, WebSocket, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from fastapi.staticfiles import StaticFiles

app = FastAPI()

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境中应该限制
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 尝试挂载静态文件，以便直接从FastAPI提供前端文件
try:
    app.mount("/", StaticFiles(directory="dist", html=True), name="static")
except Exception as e:
    print(f"挂载静态文件失败: {e}")

# 存储活跃的WebSocket连接
active_websockets = []

# 存储SSE客户端和它们的响应队列
sse_clients = {}

# 随机生成响应
async def generate_random_response(message):
    # 基础响应模板
    templates = [
        f"你好，我收到了你的消息：{message}。",
        f"正在处理你的请求：{message}。",
        f"这是一个随机回复，你发送的是：{message}。",
        f"我是一个流式响应示例，你说了：{message}。",
        f"感谢你的输入：{message}，我们正在分析中..."
    ]
    
    # 选择一个基础模板
    base_response = random.choice(templates)
    
    # 生成额外的随机文本，确保总长度约为100个字符
    additional_length = max(0, 100 - len(base_response))
    if additional_length > 0:
        # 生成随机汉字和标点
        chars = "的一是在不了有和人这中大为上个国我以要他时来用们生到作地于出就分对成会可主发年动同工也能下过子说产种面而方后多定行学法所民得经十三之进着等部度家电力里如水化高自二理起小物现实加量都两体制机当使点从业本去把性好应开它合还因由其些然前外天政四日那社义事平形相全表间样与关各重新线内数正心反你明看原又么利比或但质气第向道命此变条只没结解问意建月公无系军很情者最立代想已通并提直题党程展五果料象员革位入常文总次品式活设及管特件长求老头基资边流路级少图山统接知较将组见计别她手角期根论运农指几九区强放决西被干做必战先回则任取据处队南给色光门即保治北造百规热领七海口东导器压志世金增争济阶油思术极交受联什认六共权收证改清己美再采转更单风切打白教速花带安场身车例真务具万每目至达走积示议声报斗完类八离华名确才科张信马节话米整空元况今集温传土许步群广石记需段研界拉林律叫且究观越织装影算低持音众书布复容儿须际商非验连断深难近矿千周委素技备半办青省列习响约支般史感劳便团往酸历市克何除消构府称太准精值号率族维划选标写存候毛亲快效斯院查江型眼王按格养易置派层片始却专状育厂京识适属圆包火住调满县局照参红细引听该铁价严龙飞"
        punctuation = "，。！？；：""''（）【】《》、…—"
        
        additional_text = ''.join(random.choice(chars + punctuation) for _ in range(additional_length))
        response = base_response + additional_text
    else:
        response = base_response
    
    # 确保总长度约为100个字符
    response = response[:100]
    
    print(f"生成的响应(长度:{len(response)}): {response}")
    
    # 模拟流式返回，每0.1秒返回一个字符
    for char in response:
        yield char
        await asyncio.sleep(0.1)  # 每个字符之间的延迟改为0.1秒
    
    yield "[END]"  # 结束标记

# 只修改WebSocket相关部分
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    try:
        await websocket.accept()
        active_websockets.append(websocket)
        
        # 发送连接成功消息
        await websocket.send_text("连接已建立")
        print(f"WebSocket客户端已连接，当前活跃连接数: {len(active_websockets)}")
        
        while True:
            # 接收消息
            message = await websocket.receive_text()
            print(f"WebSocket收到消息: {message}")
            
            # 生成并发送流式响应
            async for chunk in generate_random_response(message):
                try:
                    await websocket.send_text(chunk)
                except Exception as e:
                    print(f"发送WebSocket消息时出错: {e}")
                    raise
                
    except Exception as e:
        print(f"WebSocket错误: {e}")
    finally:
        if websocket in active_websockets:
            active_websockets.remove(websocket)
            print(f"WebSocket客户端已断开，当前活跃连接数: {len(active_websockets)}")

# 只修改相关部分
@app.get("/sse")
async def sse_endpoint(request: Request):
    client_id = id(request)
    queue = asyncio.Queue()
    sse_clients[client_id] = queue
    
    async def event_generator():
        try:
            # 发送初始连接消息
            yield "data: 连接已建立\n\n"
            
            while True:
                # 等待队列中的消息或发送心跳
                try:
                    message = await asyncio.wait_for(queue.get(), timeout=30)
                    # 确保每条消息都有正确的格式：data: 消息内容\n\n
                    yield f"data: {message}\n\n"
                except asyncio.TimeoutError:
                    # 发送心跳以保持连接
                    yield "data: ping\n\n"
        except Exception as e:
            print(f"SSE事件生成器错误: {e}")
        finally:
            if client_id in sse_clients:
                del sse_clients[client_id]
                print(f"SSE客户端 {client_id} 已断开连接")
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Content-Type": "text/event-stream",
            "X-Accel-Buffering": "no"  # 禁用Nginx缓冲（如果使用Nginx）
        }
    )

@app.post("/send")
async def send_message(request: Request):
    try:
        data = await request.json()
        message = data.get("message", "")
        print(f"收到POST消息: {message}")
        
        # 创建一个响应任务
        asyncio.create_task(send_sse_response(message))
        
        return {"status": "success"}
    except Exception as e:
        print(f"处理POST请求错误: {e}")
        return {"status": "error", "message": str(e)}

async def send_sse_response(message):
    # 为所有SSE客户端生成响应
    async for chunk in generate_random_response(message):
        # 向所有SSE客户端广播消息
        for client_id, queue in list(sse_clients.items()):
            try:
                print(f"向SSE客户端 {client_id} 发送: {chunk}")
                await queue.put(chunk)
            except Exception as e:
                print(f"向SSE客户端 {client_id} 发送消息失败: {e}")

# 添加一个简单的健康检查端点
@app.get("/health")
async def health_check():
    return {"status": "ok"}

if __name__ == "__main__":
    print("启动服务器在 http://localhost:8000")
    uvicorn.run(app, host="0.0.0.0", port=8000)